package com.zeratul.common.core.context;

import com.zeratul.common.core.context.base.ZtInheritableThreadLocal;
import com.zeratul.common.core.context.enums.DataContextType;

import java.util.HashMap;
import java.util.Map;

/**
 * 数据上下文处理类
 */
public class DataContextHolde {

    /**
     * 请求上下文数据存储集合
     */
    private static final ZtInheritableThreadLocal<Map<String, Object>> REQ_CONTEXT_HOLD = new ZtInheritableThreadLocal<>();

    /**
     * 链路上下文数据存储集合
     */
    private static final ZtInheritableThreadLocal<Map<String, Object>> TRACE_CONTEXT_HOLD = new ZtInheritableThreadLocal<>();

    /**
     * 设置请求上下文数据
     */
    public static void setReqData(String key, Object value) {
        setData(DataContextType.REQUEST, key, value);
    }

    /**
     * 设置链路上下文数据集合
     */
    public static void setTraceData(String key, Object value) {
        setData(DataContextType.TRACE, key, value);
    }

    /**
     * 获取请求上下文数据
     */
    public static <R> R getReqData(String key) {
        return getData(DataContextType.REQUEST, key, null);
    }

    /**
     * 获取请求上下文的Map集合
     */
    public static Map<String, Object> getReqContextOnCopy() {
        return new HashMap<>(REQ_CONTEXT_HOLD.get(HashMap::new));
    }

    /**
     * 获取链路上下文的Map集合
     */
    public static Map<String, Object> getTraceContextOnCopy() {
        return new HashMap<>(TRACE_CONTEXT_HOLD.get(HashMap::new));
    }

    /**
     * 设置请求上下文的Map集合
     */
    public static void setReqContext(Map<String, Object> context) {
        REQ_CONTEXT_HOLD.set(context);
    }

    /**
     * 设置链路上下文的Map集合
     */
    public static void setTraceContext(Map<String, Object> context) {
        TRACE_CONTEXT_HOLD.set(context);
    }

    /**
     * 获取请求上下文数据 带默认值
     */
    public static <R> R getReqData(String key, R defaultValue) {
        return getData(DataContextType.REQUEST, key, defaultValue);
    }

    /**
     * 获取链路上下文数据
     */
    public static <R> R getTraceData(String key) {
        return getData(DataContextType.TRACE, key, null);
    }

    /**
     * 获取链路上下文数据 带默认值
     */
    public static <R> R getTraceData(String key, R defaultValue) {
        return getData(DataContextType.TRACE, key, defaultValue);
    }

    /**
     * 移除请求上下文数据
     */
    public static void removeReqData(String key) {
        remove(DataContextType.REQUEST, key);
    }

    /**
     * 移除链路上下文数据
     */
    public static void removeTraceData(String key) {
        remove(DataContextType.TRACE, key);
    }

    /**
     * 设置上下文数据，根据数据类型
     */
    private static void setData(DataContextType type, String key, Object value) {
        if (type == DataContextType.REQUEST) {
            REQ_CONTEXT_HOLD.get(HashMap::new).put(key, value);
        } else {
            TRACE_CONTEXT_HOLD.get(HashMap::new).put(key, value);
        }
    }

    /**
     * 获取上下文数据
     */
    @SuppressWarnings("unchecked")
    private static <R> R getData(DataContextType type, String key, R defaultValue) {
        if (type == DataContextType.REQUEST) {
            return (R) REQ_CONTEXT_HOLD.get(HashMap::new).getOrDefault(key, defaultValue);
        } else {
            return (R) TRACE_CONTEXT_HOLD.get(HashMap::new).getOrDefault(key, defaultValue);
        }
    }

    /**
     * 上下文数据移除
     */
    private static void remove(DataContextType type, String key) {
        if (type == DataContextType.REQUEST) {
            REQ_CONTEXT_HOLD.get(HashMap::new).remove(key);
        } else {
            TRACE_CONTEXT_HOLD.get(HashMap::new).remove(key);
        }
    }

    /**
     * 销毁上下文对象
     */
    public static void clear() {
        REQ_CONTEXT_HOLD.remove();
        TRACE_CONTEXT_HOLD.remove();
    }
}
